import tensorflow as tf
from matplotlib import pyplot as plt
from keras.datasets import mnist
from keras.preprocessing.image import ImageDataGenerator

# 加载MNIST数据集
(x_train, y_train), (x_test, y_test) = mnist.load_data()

# MNIST数据集是灰度图像，需要将其扩展为具有单通道的四维张量
x_train = x_train.reshape(-1, 28, 28, 1)
x_test = x_test.reshape(-1, 28, 28, 1)

# 将像素值标准化为[0, 1]区间
x_train = x_train / 255.0
x_test = x_test / 255.0

# 定义图像数据增强器
image_data_generator = ImageDataGenerator(rotation_range=40,
                                          width_shift_range=0.2,
                                          height_shift_range=0.2,
                                          shear_range=20,
                                          zoom_range=0.2,
                                          horizontal_flip=True,
                                          vertical_flip=False)  # MNIST图像不适合垂直翻转

# 选择一个示例图片
index = 0
original_image = x_train[index]

# 生成增强后的图像
augmented_image = image_data_generator.flow(x_train, shuffle=False).next()

# 显示原始图像和增强图像
plt.subplot(1, 2, 1)
plt.imshow(original_image.squeeze(), cmap='gray')
plt.title('Original Image')
plt.subplot(1, 2, 2)
plt.imshow(augmented_image[0].squeeze(), cmap='gray')
plt.title('Augmented Image')
plt.show()
